package m3f.factorization;

import m3f.data.MultimodalCachedLoader;
import m3f.data.MultimodalDataProvider;
import m3f.data.MultimodalDataSet;
import m3f.io.DenseMatrixWriter;
import m3f.matrix.Minibatch;

import org.jblas.DoubleMatrix;
import org.jblas.Solve;

public class OnlineJointFactorization {
	
	protected MultimodalDataProvider trainingData;
	protected DoubleMatrix P;
	protected DoubleMatrix Q;
	protected DoubleMatrix G;
	protected DoubleMatrix lambdaI;
	protected double visualError;
	protected double textError;
	protected int factors;
	protected double omega;
	
	public OnlineJointFactorization(MultimodalDataProvider set){
		trainingData = set;
	}
	
	protected void computeError(){
		System.out.println("Computing error");
		trainingData.reset(false);
		this.textError = 0;
		this.visualError = 0;
		long s = System.currentTimeMillis();
		computeProjectionOperator();
		System.out.println("Calculating projection operator: " + (System.currentTimeMillis()-s) + " ms");
		ConcurrentErrorComputation cec = new ConcurrentErrorComputation((MultimodalCachedLoader)trainingData, P, Q, G, omega);
		Thread thread = new Thread(cec);
		thread.start();
		try{
			thread.join();
		}catch(InterruptedException ie){
			ie.printStackTrace();
		}
		this.textError = cec.getTextError();
		this.visualError = cec.getVisualError();
	}
	
	protected void computeProjectionOperator(){
		DoubleMatrix A;
		G = P.transpose().mmul(P).muli(1 - omega);
		A = Q.transpose().mmul(Q).muli(omega);
		G.addi(A).addi(lambdaI);
	}
	
	protected void initialize(int factors, int visualFeatures, int textFeatures){
		System.out.println("Initializing model with " + visualFeatures + " visual features, " +
				textFeatures + " text features and " + factors + " latent factors.");
		System.out.println("Training instances: " + this.trainingData.matrixRows());
		P = DoubleMatrix.rand(visualFeatures, factors).muli(0.01);
		Q = DoubleMatrix.rand(textFeatures, factors).muli(0.01);
	}
	
	public static DoubleMatrix latentRepresentation(Minibatch V, Minibatch T, DoubleMatrix P, DoubleMatrix Q, DoubleMatrix G, double omega){
		DoubleMatrix A, B, H;
		A = P.transpose().mmul(V.matrix()).muli(1 - omega);
		B = Q.transpose().mmul(T.matrix()).muli(omega);
		A.addi(B);
		H = Solve.solve(G, A);
		return H;
	}
	
	public DoubleMatrix projectVisualData(MultimodalDataSet test, double xi){
		System.out.println("Test set with " + test.matrixRows() + " samples with " + test.visualFeatures() + " visual features and " + test.textFeatures() + " text features");
		Minibatch V = test.maxVisualBatch();
		Minibatch T = test.maxTextBatch();
		DoubleMatrix A,U,Y;
		A = Q.transpose().mmul(Q);
		A.muli(xi).addi( P.transpose().mmul(P) );
		//A = P.transpose().mmul(P);
		U = Solve.solve(A, P.transpose().mmul(V.matrix()));
		Y = Q.mmul(U);
		Y.subi(T.matrix());
		double error = Math.pow(Y.norm2(),2.0);
		System.out.println("Coding error: " + error);
		return U;
	}
	
	public DoubleMatrix projectTextData(MultimodalDataSet test, double xi){
		System.err.println("THIS METHOD IS NOT IMPLEMENTED YET!");
		throw new RuntimeException();/*
		System.out.println("Test set with " + test.matrixRows() + " samples with " + test.visualFeatures() + " visual features and " + test.textFeatures() + " text features");
		Minibatch V = test.maxVisualBatch();
		Minibatch T = test.maxTextBatch();
		DoubleMatrix A,U,Y;
		A = Q.transpose().mmul(Q);
		A.muli(xi).addi( P.transpose().mmul(P) );
		U = Solve.solve(A, P.transpose().mmul(V.matrix()));
		Y = Q.mmul(U);
		Y.subi(T.matrix());
		double error = Math.pow(Y.norm2(),2.0);
		System.out.println("Coding error: " + error);
		return U;		*/
	}
	
	public void run(int latentFactors, double lambda, double gamma0, double alpha, double beta, double omega, int epochs, int minibatch, boolean skipError){
		this.factors = latentFactors;
		this.omega = omega;
		this.lambdaI = DoubleMatrix.eye(factors).muli(lambda);
		trainingData.setMinibatchSize(minibatch);
		initialize(factors, trainingData.visualFeatures(), trainingData.textFeatures());
		if(!skipError){
			computeError();
			System.out.println("Initial error: visualError=" + visualError + " textError=" + textError);
		}
		int t = 1;
		DoubleMatrix A, B, C, hh;
		for(int i = 0; i < epochs; i++){
			trainingData.reset(true);
			System.out.println("Training epoch " + i);
			while(trainingData.hasMoreData()){
				long allStar = System.currentTimeMillis();
				long star = System.currentTimeMillis();
				double gamma = gamma0/(1 + gamma0*lambda*t);
				// Prepare minibatch
				Minibatch vv = trainingData.nextVisual();
				Minibatch tt = trainingData.nextTextual();
				// 1. Compute the latent representation for the minibatch:
				computeProjectionOperator();
				hh = OnlineJointFactorization.latentRepresentation(vv, tt, P, Q, G, omega);
				System.out.print("Latent representation (" + (System.currentTimeMillis()-star) + "ms) ");
				star = System.currentTimeMillis();
				// 2. Update model matrices:
				double omga = 1 - gamma*alpha;
				double omgb = 1 - gamma*beta;
				C = hh.mmul(hh.transpose());
				System.out.print("hh' (" + (System.currentTimeMillis()-star) + "ms) ");
				star = System.currentTimeMillis();
				// Update P
				A = vv.matrix().mmul(hh.transpose()).mul(gamma*(1-omega));
				B = P.mmul(C).mul(-gamma*(1-omega));
				P.muli(omga).addi(A).addi(B);
				System.out.print("Updating P (" + (System.currentTimeMillis()-star) + "ms) ");
				star = System.currentTimeMillis();
				// Update Q
				A = tt.matrix().mmul(hh.transpose()).mul(gamma*omega);
				B = Q.mmul(C).mul(-gamma*omega);
				Q.muli(omgb).addi(A).addi(B);
				System.out.print("Updating Q (" + (System.currentTimeMillis()-star) + "ms) ");
				star = System.currentTimeMillis();
				t++;
				System.out.println(" All minibatch = " + (System.currentTimeMillis()-allStar) + "ms");
				star = System.currentTimeMillis();
			}
			System.out.println("Epoch " + i + " done");
			if(!skipError){
				computeError();
				System.out.println("Epoch " + i + ": visualError=" + visualError + " textError=" + textError);
			}
		}
	}
	
	public void writeLatentRepresentation(String outputFile){
		System.out.println("Writing the latent representation for training data");
		trainingData.reset(false);
		DenseMatrixWriter writer = new DenseMatrixWriter(outputFile);
		computeProjectionOperator();
		ConcurrentMatrixWriter cmw = new ConcurrentMatrixWriter(trainingData, writer, P, Q, G, omega);
		Thread thread = new Thread(cmw);
		thread.start();
		try{
			thread.join();
		}catch(InterruptedException ie){
			ie.printStackTrace();
		}
		writer.close();
	}
	
	public DoubleMatrix getP() {
		return P;
	}
	
	public DoubleMatrix getQ(){
		return Q;
	}
}